'''
Created on Sep 29, 2009

@author: mkiyer
'''

from veggie.gene.gene import parse_refseq_genes
from veggie.sample.samplegroup import parse_samplegroups_xml
from optparse import OptionParser
import h5py
import matplotlib.pyplot as plt
import numpy as np
import veggie.coverage.covdb as covdb
import operator
import os
import sys

def get_exon_coverage(gene, dset):
    exonlength = 0
    exoncov = 0
    for e in gene.exons:
        exoncov += sum(dset[gene.chrom][e[0]:e[1]])
        exonlength += e[1] - e[0]
    return 1000.0 * exoncov / exonlength

def plot_scatter(x, y, name):
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.scatter(x,y)
    ax.grid(True)
    plt.savefig('%s.png' % name)
    plt.show()
    
if __name__ == '__main__':
    optionparser = OptionParser("usage: %prog [options]")
    optionparser.add_option("-1", dest="unique")
    optionparser.add_option("-2", dest="duplicate")
    optionparser.add_option("-s", "--samples", dest="samples",
                            help="sample XML file to process")
    optionparser.add_option("-o", "--output-dir", dest="output_dir",
                            default=os.path.curdir,
                            help="path to place the output files")
    (options, args) = optionparser.parse_args()


    # open the coverage database file
    h5unique = h5py.File(options.unique, 'r')
    h5dup = h5py.File(options.duplicate, 'r')
    # get genes
    genes = list(parse_refseq_genes())

    for unique_name in h5unique:
        
        if unique_name in h5dup:

            dset_unique = h5unique[unique_name]
            dset_dup = h5dup[unique_name]

            total_unique = dset_unique.attrs['read_total']
            total = dset_dup.attrs['read_total']
            
            print 'comparing', unique_name, '...'
            print 'unique:', total_unique
            print 'total:', total
            print '%unique:', float(total_unique) / total
            
            cov_unique = np.zeros(len(genes), dtype=np.float)
            cov_dup = np.zeros(len(genes), dtype=np.float)
            gene_cov = []
            
            for i, g in enumerate(genes):
                if i % 1000 == 0:
                    print i
                cov_unique[i] = get_exon_coverage(g, dset_unique)
                cov_dup[i] = get_exon_coverage(g, dset_dup)
                gene_cov.append((g, cov_unique[i], cov_dup[i]))

            gene_cov = sorted(gene_cov, key=operator.itemgetter(2))
                
            f = open('%s_unique_vs_dup.txt' % unique_name, 'w')
            for g, u, d in gene_cov:
                f.write('%s\t%s\t%f\t%f\n' % 
                        (g.symbol, g.acc, u, d))

            plot_scatter(cov_unique, cov_dup, unique_name)
            sys.exit(0)
